# regTable.R

# Takes a list of regression models.  Returns a matrix in which the columns 
# are estimates and standard errors -- two columns for each model.  Output is 
# to the screen; there is no LaTeX formatting.

regTable <- function (
  objList, 
  colNames     = NULL, 
  rowsToRemove = NULL, 
  rowsToKeep   = NULL,
  clusterSEs   = FALSE,
  clusterVar   = NULL) {

  require(dplyr)         # for %>%
  require(lmtest)        # for coeftest()
  require(multiwayvcov)  # for cluster.vcov, which helps to compute clustered SEs for OLS estimates 
  
  if (! class(objList) %in% 'list') {
    stop("objList must be of class 'list'.")
  }
  if (! is.null(colNames) && length(colNames) != length(objList)) {
    stop("colNames must be NULL or the length of objList.")
  }
  classVec <- sapply(objList, class)  
  if (is.matrix(classVec)) {
    classVec <- classVec[1,]  # get first class of all models
  }
  if (! all(classVec %in% c('lm', 'plm', 'ivreg'))) {
    warning("regTable() has only been designed to work with models of class 'lm', 'plm', and 'ivreg'.")
  }
  if (clusterSEs & !all(classVec == 'lm') & !all(classVec == 'ivreg')) {
    stop("clusterSEs is TRUE, but regTable() can only cluster SEs only when all objects in objList are 'lm' objects or when all objects in objList are 'ivreg' objects.")
  }
  if (!is.null(clusterVar) && class(clusterVar) != 'list') {
    stop("clusterVar must be an object of class 'list'")
  }
  
  
  # Get column names
  if (is.null(colNames)) {
    tmpNames <- NULL
    for (obj in objList) {
      if (any(c('iv_robust', 'lm_robust') %in% class(obj))) {  # objects from the "estimatr" package
        tmpNames <- c(tmpNames, obj$outcome)
      }
      if ('ivreg' %in% class(obj)) {
        tmpNames <- c(tmpNames, obj$formula[2])
      }
      else if ('lm' %in% class(obj)) {
        tmpNames <- c(tmpNames, as.character(obj$terms[[2]])[1])
      }
    }
    colNames <- tmpNames
  }
  if (! is.null(colNames)) {
    colNames <- paste0(rep(colNames, each = 2), c('', '.se'))
  }

  
  
  # Get names of predictors, including intercept.  Eliminate names of 
  # unwanted rows.
  tmp      <- sapply(objList, function (x) { names(x$coefficients) } )
  rowNames <- unique(unlist(tmp)) 
  if (! is.null(rowsToRemove)) { 
    for (pat in rowsToRemove) {
      rowNames <- rowNames[!grepl(pat, rowNames)]
    }
  }
  else if (! is.null(rowsToKeep)) {
    rowsToKeep <- paste0(rowsToKeep, collapse = '|')
    keepRegex  <- paste0('^(', rowsToKeep, ')')
    rowNames <- rowNames[grepl(keepRegex, rowNames)]      
  }
  

  
  # Set up output matrix
  out      <- matrix(
    nrow     = length(rowNames), 
    ncol     = length(objList)*2, 
    dimnames = list(rowNames, colNames))
  
  
  # Get clustered vcov matrices for OLS estimates.  Then fill the output 
  # matrix that this function will return.  Use coeftest(x) here because it is 
  # more robust than summary(x)$coefficients.  
  if (clusterSEs == TRUE & all(classVec == 'lm')) {
    vcovs.clustered <- mapply(
      FUN     = cluster.vcov,
      model   = objList,
      cluster = clusterVar)
    coefsAndSEs <- mapply(
      FUN   = coeftest,
      x     = objList,
      vcov. = vcovs.clustered) 
    coefsAndSEs <- sapply(coefsAndSEs, function (x) x[, c('Estimate', 'Std. Error')])
  }
  else if (clusterSEs == TRUE & all(classVec == 'ivreg')) {
    coefsAndSEs <- mapply(
      FUN = cluster.robust.se,
      objList,
      clusterVar)
    coefsAndSEs <- lapply(coefsAndSEs, function (x) x[, c('Estimate', 'Std. Error')])    
  }
  else {  
    coefsAndSEs <- lapply(objList, function (x) {
      tmp <- coeftest(x)[, c('Estimate', 'Std. Error')]

      # If we are dealing with a single predictor, the line above will drop 
      # the row name, causing problems. This block of code restores the row
      # name.  
      if (class(tmp) == 'numeric') {  # if we're dealing with a single row of output
        tmp <- matrix(
          data = tmp,
          dimnames = list(c('Estimate', 'Std. Error'), rownames(coeftest(x)))) %>%
          t
      }
      tmp
    })
  }
  
  for (i in 1:length(coefsAndSEs)) {
    rowNamesToUse <- rowNames[rowNames %in% rownames(coefsAndSEs[[i]])]
    out[rowNamesToUse, (i*2 -1):(i*2)] <- coefsAndSEs[[i]][rowNamesToUse, ]
  } 

  class(out) <- c('regTable', class(out))  
  out
}


# Set up print method for regTable, so that print(regTable) doesn't list 
# attributes at the bottom.
print.regTable <- prmatrix
